import argparse
import os.path

import torch

from gxl_ai_utils.utils import utils_file
from do_split_and_concat import do_split


def get_parser():
    """"""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--wav_scp_path",
        required=True,
        help="The path of wav.scp.",
    )
    parser.add_argument(
        "--text_scp_path",
        required=True,
        help="",
    )
    parser.add_argument(
        "--output_dir",
        required=True,
        help="",
    )
    parser.add_argument(
        "--split_num",
        required=True,
        type=int,
        help="The path of text.scp.",
    )
    parser.add_argument(
        "--num_thread_for_manifest",
        required=True,
        type=int,
        help="The path of text.scp.",
    )
    parser.add_argument(
        "--num_thread_for_fbank",
        required=True,
        type=int,
        help="The path of text.scp.",
    )
    parser.add_argument(
        "--perturb_speed",
        required=False,
        default=False,
        type=bool,
        help="",
    )
    parser.add_argument(
        "--partition",
        required=True,
        type=str,
        help="The path of text.scp.",
    )
    parser.add_argument(
        "--prefix",
        required=True,
        type=str,
        help="The path of text.scp.",
    )
    return parser


def main():
    """"""
    parser = get_parser()
    args = parser.parse_args()
    torch.set_num_threads(1)
    torch.set_num_interop_threads(1)
    wav_scp_path, wer_scp_path, output_dir, split_num = args.wav_scp_path, args.text_scp_path, args.output_dir, args.split_num
    utils_file.logging_print(f'args: {args}')
    utils_file.logging_print(f'getted task args finished')

    utils_file.logging_print(f'start split')
    output_dir_split = os.path.join(output_dir, 'split')
    utils_file.logging_print(f'output_dir_split: {output_dir_split}')
    utils_file.makedir_sil(output_dir_split)
    do_split(wav_scp_path, wer_scp_path, output_dir_split, split_num=split_num)
    utils_file.logging_print(f'split finished')

    utils_file.logging_print('start handle per split')
    for i in range(split_num):
        utils_file.logging_print(f'start handle for split_{i}')
        root_dir_i = os.path.join(output_dir_split, f'temp_{i}')
        manifest_dir_i = os.path.join(root_dir_i, 'manifest')
        fbank_dir_i = os.path.join(root_dir_i, 'fbank')
        utils_file.makedir_sil(manifest_dir_i)
        utils_file.makedir_sil(fbank_dir_i)
        utils_file.logging_print(f'manifest_dir_i: {manifest_dir_i}')
        utils_file.logging_print(f'fbank_dir_i: {fbank_dir_i}')

        utils_file.logging_print(f'start handle manifest for split_{i}')
        now = utils_file.do_get_now_time()
        wav_path_i = os.path.join(root_dir_i, 'wav.scp')
        text_path_i = os.path.join(root_dir_i, 'text')
        wav_mani_name, text_mani_name = utils_file.get_jsonl_filename4icefall(args.prefix, args.partition)
        # 会自动跳过处理过的
        utils_file._do_convert_scp_to_manifest4icefall(
            wav_scp_path=wav_path_i,
            text_scp_path=text_path_i,
            wav_manifest_path=os.path.join(manifest_dir_i, wav_mani_name),
            text_manifest_path=os.path.join(manifest_dir_i, text_mani_name),
            num_thread=args.num_thread_for_manifest
        )
        utils_file.logging_print(f'finish handle manifest for split_{i}, time: {utils_file.do_get_now_time() - now}s')


        utils_file.logging_print(f'start handle fbank for split_{i}')
        now = utils_file.do_get_now_time()
        # 会自动跳过处理过的
        utils_file._do_compute_fbank4icefall(
            perturb_speed=args.perturb_speed,
            manifests_dir=manifest_dir_i,
            fbank_dir=fbank_dir_i,
            partition=args.partition,
            prefix=args.prefix,
            num_jobs=args.num_thread_for_fbank,
        )
        utils_file.logging_print(f'finish handle fbank for split_{i}, time: {utils_file.do_get_now_time() - now}s')

    utils_file.logging_print(f'start concat')
    now = utils_file.do_get_now_time()
    output_dir_concat = output_dir
    res_cuts_path = os.path.join(output_dir_concat, f'{args.prefix}_cuts_{args.partition}.jsonl.gz')
    res_list = []
    for i in range(split_num):
        utils_file.logging_print(f'start handle for split_{i}')
        root_dir_i = os.path.join(output_dir_split, f'temp_{i}')
        fbank_dir_i = os.path.join(root_dir_i, 'fbank')
        cuts_file_path_i = os.path.join(fbank_dir_i, f'{args.prefix}_cuts_{args.partition}.jsonl.gz')
        res_list.extend(utils_file.load_list_file_clean(cuts_file_path_i))
    utils_file.write_list_to_file(res_list, res_cuts_path)
    utils_file.logging_print(f'concat finished,耗时{utils_file.do_get_now_time() - now}s')


if __name__ == "__main__":
    main()
